
import requests
import json
import re
import pandas as pd
import time
import os
import sys

class wenxin_api(object):

    def __init__(self) -> None:
        pass
        
    def get_access_token(self):
        """
        使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
        """
        url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=DLJhqAcvQlhxzgdqcLU3JmhI&client_secret=9jAXdLnPkfT5uZ5zmB0V7V2kyGOwnMUt"    
        
        payload = json.dumps("")
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        return response.json().get("access_token")


    def _chat_with_ernie_speed_128K(self, user_prompt):
        
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token=" + self.get_access_token()
        payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": user_prompt
            }
        ],
        "temperature": 0.9
        })
        headers = {
            'Content-Type': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        # print(response.text)
        return response.text

tag_template = '''假设我是一名图书编辑，需要为一系列图书打上准确的标签。请用一段话描述“{}”这一图书标签的核心意义。'''

def pro_tag(tag_name):
    api = wenxin_api()
    gpt_ans = api._chat_with_ernie_speed_128K(tag_template.format(tag_name))                                        
    gpt_ans = json.loads(gpt_ans) 
    output = gpt_ans['result']
    return output

def main(input_file='../Dataset/tag_data/标签库.txt', output_file='./processed_data/SourceUniversity/proces_tongxueshe_tag.csv'):
    """主函数，读取标签库文件并输出到CSV文件"""

    txt = ""
    # 打开文件，准备读取  
    with open(input_file, 'r', encoding='utf-8') as file:  # 使用utf-8编码以支持多语言文本  
        # 逐行读取文件  
        for line in file:  
            # 去除每行末尾的换行符，并计算字符长度  
            stripped_line = line.strip()  
            if len(stripped_line) >= 5:   
                txt += stripped_line + "、"

    tag_lt = [t_s for t_s in txt.split('、') if t_s != '']
    print(f"len(tag_lt) : {len(tag_lt)}")
    
    
    info_lt = []  # 存储处理结果的列表
    index = 0  # 从哪个索引开始处理

    directory = os.path.dirname(output_file)
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"目录 {directory} 不存在，已创建。")

    if os.path.exists(output_file):
        # 读取已有的数据，确定从哪个索引开始
        df = pd.read_csv(output_file, encoding='utf-8-sig')
        # if 'tag_name' in df.columns:
        index = len(df)+1
    else:
        with open(output_file, 'w', encoding='utf-8') as file:
            print("创建文件来存储标签的描述信息，结果存储在："+output_file)
            import csv
            writer = csv.writer(file)
            writer.writerow(['tag_name', 'tag_describe'])
            pass 
    # 从指定索引开始处理
    for t_name in tag_lt[index:]:
        p_tag = pro_tag(t_name)
        info_lt.append([t_name, p_tag])
        time.sleep(20)  # 模拟耗时操作
        print("generate: ", index, "   tag: ", t_name)
        index += 1

        # 将生成的数据写入CSV文件
        df = pd.DataFrame(info_lt, columns=['tag_name', 'tag_describe'])
        df.to_csv(output_file, index=False, mode='a', encoding='utf-8-sig', header=False)
        info_lt = []  # 清空列表，准备下一次写入
    print(f"标签信息已成功写入到 {output_file}")

if __name__ == "__main__":
    import json
    if len(sys.argv) != 2:
        print("使用方法: python gener_tag_info.py [配置文件路径]")
        sys.exit(1)
    config_path = sys.argv[1]
    with open(config_path, 'r') as f:
        config = json.load(f)

    tag_file_path = config['tag_path']
    tag_info_path = config['tag_info_path']
    main(tag_file_path, tag_info_path)









